# -*- coding: utf-8 -*-
# 模块导入
import numpy as np
import netCDF4 as nc
from osgeo import gdal, osr, ogr
import os
from tqdm import tqdm
import glob


# 单个nc数据ndvi数据读取为多个tif文件，并将ndvi值化为-1-1之间
def NC_to_tiffs(NC_file: str, out_filename: str, output_folder: str,
                start_year: int,
                key: str,
                reverse=False,
                factor=1):
    nc_file = nc.Dataset(NC_file)
    Lon = nc_file.variables['lon'][:]
    Lat = nc_file.variables['lat'][:]
    arr = np.asarray(nc_file.variables[key])  # 将数据读取为数组
    arr_float = arr.astype(float) / factor  #将int类型改为float类型,并化为-1 - 1之间
    # 影像的左上角和右下角坐标
    LonMin, LatMax, LonMax, LatMin = [Lon.min(), Lat.max(), Lon.max(), Lat.min()]
    # 分辨率计算
    N_Lat = len(Lat)
    N_Lon = len(Lon)
    Lon_Res = (LonMax - LonMin) / (float(N_Lon) - 1)
    Lat_Res = (LatMax - LatMin) / (float(N_Lat) - 1)
    for i in tqdm(range(len(arr[:]))):  # 0维大小
        # 创建.tif文件
        driver = gdal.GetDriverByName('GTiff')
        out_tif_name = output_folder + '\\' + out_filename + '_' + str(i + start_year) + '.tif'
        out_tif = driver.Create(out_tif_name, N_Lon, N_Lat, 1, gdal.GDT_Float32)
        # 设置影像的显示范围
        # -Lat_Res一定要是-的
        geotransform = (LonMin, Lon_Res, 0, LatMax, 0, -Lat_Res)
        out_tif.SetGeoTransform(geotransform)
        # 获取地理坐标系统信息，用于选取需要的地理坐标系统
        srs = osr.SpatialReference()
        srs.ImportFromEPSG(4326)  # 定义输出的坐标系为"WGS 84"，AUTHORITY["EPSG","4326"]
        out_tif.SetProjection(srs.ExportToWkt())  # 给新建图层赋予投影信息
        # 数据写出
        if reverse:  # 转置, 有些NC文件中数据是反的，需要转置
            tarr = arr_float[i][::].T
        else:
            tarr = arr_float[i][::]
        out_tif.GetRasterBand(1).WriteArray(tarr[::-1])  # 将数据写入内存，此时没有写入硬盘 此处[::-1]用于图像的垂直镜像对称，避免图像颠倒
        out_tif.FlushCache()  # 将数据写入硬盘
        out_tif = None  # 注意必须关闭tif文件


if __name__ == '__main__':

    """
    # 该文件经纬度是反的, 所以要转置
    sig_05deg_yearly_1981_2022.nc   dict_keys(['sig', 'lon', 'lat'])            (42, 720, 360)
    wy_0p5deg_yearly_1981_2022.nc   dict_keys(['wy', 'lon', 'lat', 'time'])     (42, 720, 360)
    wye_05deg_yearly_1981_2022.nc   dict_keys(['wye', 'lon', 'lat', 'time'])    (42, 720, 360)
    """

    input_file = r'C:\Users\calibre66\Desktop\WaterYield1981-2022\wye_05deg_yearly_1981_2022.nc'
    output_folder = r'C:\Users\calibre66\Desktop\WaterYield1981-2022\tiffs'
    output_tif = 'wye_05deg'
    key = 'wye'  # dict_keys=['wy', 'lon', 'lat', 'time']
    start_year = 1981
    end_year = 2022

    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    NC_to_tiffs(input_file, output_tif, output_folder,
                start_year,  key, True, 1)
    # 这套数据行和列是反的，所以要转置

    # # 文件夹批量读取nc文件
    # Input_folder = r'C:\Users\calibre66\Desktop\WaterYield1981-2022'
    # Output_folder = r'C:\Users\calibre66\Desktop\WaterYield1981-2022\tiffs'
    # output_filenames = ["sig_05deg_yearly_1981_2022", "wy_0p5deg_yearly_1981_2022", "wye_05deg_yearly_1981_2022"]
    # if not os.path.exists(Output_folder):
    #     os.makedirs(Output_folder)
    #
    # start_year = 1981
    # end_year = 2022
    # # dict_keys=['wy', 'lon', 'lat', 'time']
    # key = 'wy'
    # # 读取所有nc数据
    # data_list = glob.glob(Input_folder + '\\*.nc')
    # for i in range(len(data_list)):
    #     data = data_list[i]
    #     NC_to_tiffs(data, Output_folder)
    #     print(data + '-----转tif成功')
    # print('----转换结束----')
